#include <asm.h>
#include <csr.h>

.macro SAVE_CONTEXT 
  
  /* TODO: save all general purpose registers here! */
  ld t1,0(sp)
  addi sp,sp,-8
  sd a0,0(sp)
  csrr a0, CSR_MHARTID
  slli a0,a0,3
  la tp,current_running
  add tp,tp,a0
  ld tp,0(tp)
  add t1,t1,tp
  ld a0,0(sp)
  addi sp,sp,8

  sd zero,OFFSET_REG_ZERO(t1)
  //return address
  sd ra,OFFSET_REG_RA(t1)
  //pointers
  addi sp,sp,16
  sd sp,OFFSET_REG_SP(t1)
  addi sp,sp,-16
  sd gp,OFFSET_REG_GP(t1)
  sd tp,OFFSET_REG_TP(t1)
  //temporary
  
  sd t0,OFFSET_REG_T0(t1)
  sd t2,OFFSET_REG_T2(t1)
  //saved registers
  sd s0,OFFSET_REG_S0(t1)
  sd s1,OFFSET_REG_S1(t1)
  //args
  sd a0,OFFSET_REG_A0(t1)
  sd a1,OFFSET_REG_A1(t1)
  sd a2,OFFSET_REG_A2(t1)
  sd a3,OFFSET_REG_A3(t1)
  sd a4,OFFSET_REG_A4(t1)
  sd a5,OFFSET_REG_A5(t1)
  sd a6,OFFSET_REG_A6(t1)
  sd a7,OFFSET_REG_A7(t1)
  //saved registers
  sd s2,OFFSET_REG_S2(t1)
  sd s3,OFFSET_REG_S3(t1)
  sd s4,OFFSET_REG_S4(t1)
  sd s5,OFFSET_REG_S5(t1)
  sd s6,OFFSET_REG_S6(t1)
  sd s7,OFFSET_REG_S7(t1)
  sd s8,OFFSET_REG_S8(t1)
  sd s9,OFFSET_REG_S9(t1)
  sd s10,OFFSET_REG_S10(t1)
  sd s11,OFFSET_REG_S11(t1)
  //temporary 
  sd t3,OFFSET_REG_T3(t1)
  sd t4,OFFSET_REG_T4(t1)
  sd t5,OFFSET_REG_T5(t1)
  sd t6,OFFSET_REG_T6(t1)
  
  /*
   * Disable user-mode memory access as it should only be set in the
   * actual user copy routines.
   *
   * Disable the FPU to detect illegal usage of floating point in kernel
   * space.
   */
  li t0, SR_SUM | SR_FS

  /* TODO: save sstatus, sepc, stval, scause and sscratch on user stack */
  //privileged
  csrr t3,sstatus
  sd t3,OFFSET_REG_SSTATUS(t1)
  csrr t3,sepc
  sd t3,OFFSET_REG_SEPC(t1)
  
  csrr t3,stval
  sd t3,OFFSET_REG_SBADADDR(t1)
  csrr t3,scause
  sd t3,OFFSET_REG_SCAUSE(t1)
  
  ld t3,8(sp)
  sd t3,OFFSET_REG_T1(t1)
  ld t3,OFFSET_REG_T3(t1)

  csrr a0, CSR_MHARTID
  slli a0,a0,3
  la t1,current_running
  add t1,t1,a0
  ld  tp,0(t1)
  addi sp,sp,16
.endm

.macro RESTORE_CONTEXT 
  /* TODO: Restore all registers and sepc,sstatus */
  ld t1,0(sp)
  add t1,tp,t1
  addi sp,sp,16
  /*addi sp,sp,-16
  sd a0,0(sp)
  csrr a0, CSR_MHARTID
  slli a0,a0,3
  add t1,t1,a0
  ld a0,0(sp)
  addi sp,sp,16*/
  ld zero,OFFSET_REG_ZERO(t1)
  //return address
  ld ra,OFFSET_REG_RA(t1)
  //pointers
  ld sp,OFFSET_REG_SP(t1)
  ld gp,OFFSET_REG_GP(t1)
  //ld tp,OFFSET_REG_TP(t1)
  //temporary
  
  ld t0,OFFSET_REG_T0(t1)
  ld t2,OFFSET_REG_T2(t1)
  //saved registers
  ld s0,OFFSET_REG_S0(t1)
  ld s1,OFFSET_REG_S1(t1)
  //args
  ld a0,OFFSET_REG_A0(t1)
  ld a1,OFFSET_REG_A1(t1)
  ld a2,OFFSET_REG_A2(t1)
  ld a3,OFFSET_REG_A3(t1)
  ld a4,OFFSET_REG_A4(t1)
  ld a5,OFFSET_REG_A5(t1)
  ld a6,OFFSET_REG_A6(t1)
  ld a7,OFFSET_REG_A7(t1)
  //saved registers
  ld s2,OFFSET_REG_S2(t1)
  ld s3,OFFSET_REG_S3(t1)
  ld s4,OFFSET_REG_S4(t1)
  ld s5,OFFSET_REG_S5(t1)
  ld s6,OFFSET_REG_S6(t1)
  ld s7,OFFSET_REG_S7(t1)
  ld s8,OFFSET_REG_S8(t1)
  ld s9,OFFSET_REG_S9(t1)
  ld s10,OFFSET_REG_S10(t1)
  ld s11,OFFSET_REG_S11(t1)
  //temporary 
  ld t3,OFFSET_REG_T3(t1)
  ld t4,OFFSET_REG_T4(t1)
  ld t5,OFFSET_REG_T5(t1)
  ld t6,OFFSET_REG_T6(t1)

  ld t3,OFFSET_REG_SSTATUS(t1)
  csrw sstatus,t3
  ld t3,OFFSET_REG_SEPC(t1)
  csrw sepc,t3
  ld t3,OFFSET_REG_SBADADDR(t1)
  csrw stval,t3
  ld t3,OFFSET_REG_SCAUSE(t1)
  csrw scause,t3

  ld t3,OFFSET_REG_T3(t1)
  ld t1,OFFSET_REG_T1(t1)
.endm

ENTRY(lock_kernel)
    la      t1,kernel_lock_key
1:
    lr.d    t0,0(t1)
    bnez    t0,1b
    li      t0,1
    sc.d    t0,t0,0(t1)
    bnez    t0,1b
    jr ra
ENDPROC(lock_kernel)

ENTRY(unlock_kernel)
    la      t1,kernel_lock_key
2:
    lr.d    t0,0(t1)
    li      t0,0
    sc.d    t0,t0,0(t1)
    bnez    t0,2b
    jr ra
ENDPROC(unlock_kernel)
ENTRY(enable_preempt)
  /*ld t1, current_running
  ld t0, PCB_PREEMPT_COUNT(t1)
  beq t0, zero, do_enable
  addi t0, t0, -1
  sd t0, PCB_PREEMPT_COUNT(t1)
  beq t0, zero, do_enable*/
  jr ra
do_enable:
  not t0, x0
  csrs CSR_SIE, t0
  jr ra
ENDPROC(enable_preempt)

ENTRY(disable_preempt)
  /*csrw CSR_SIE, zero
  ld t1, current_running
  ld t0, PCB_PREEMPT_COUNT(t1)
  addi t0, t0, 1
  sd t0, PCB_PREEMPT_COUNT(t1)*/
  jr ra
ENDPROC(disable_preempt)

ENTRY(enable_interrupt)
  //li t0, SR_SIE
  //csrw CSR_SSTATUS, t0
  jr ra
ENDPROC(enable_interrupt)

ENTRY(disable_interrupt)
  //li t0, SR_SIE
  //csrw CSR_SSTATUS, t0
  jr ra
ENDPROC(disable_interrupt)

ENTRY(exception_handler_entry)
  addi sp,sp,-16
  li tp,OFFSET_SIZE
  sd tp,0(sp)
  sd t1,8(sp)
  SAVE_CONTEXT
  csrw CSR_SSCRATCH, x0
  
  .option push
  .option norelax
  la gp, __global_pointer$
  .option pop
  addi sp,sp,-24
  sd ra,0(sp)
  sd t0,8(sp)
  sd t1,16(sp)
  jal lock_kernel
  ld ra,0(sp)
  ld t0,8(sp)
  ld t1,16(sp)
  addi sp,sp,24
  la t0,current_cpu
  csrr a0, CSR_MHARTID
  sd a0,0(t0)
  addi t0,tp,576
  ld sp,0(t0)
  mv s0,sp
  //lui ra,%hi(ret_from_exception)
  //addi ra,ra,%lo(ret_from_exception)
  la ra,ret_from_exception
  li a0,OFFSET_SIZE
  add a0,a0,tp
  csrr a2, scause
  csrr a1, CSR_STVAL
 
  j interrupt_helper
ENDPROC(exception_handler_entry)

ENTRY(do_scheduler)
  addi sp,sp,-16
  sd zero,0(sp)
  sd t1,8(sp)
  SAVE_CONTEXT

  jal scheduler

  la tp,current_running
  ld  tp,0(tp)

  csrr a0, CSR_MHARTID
  slli a0,a0,3
  la tp,current_running
  add tp,tp,a0
  ld  tp,0(tp)

  addi sp,sp,-16
  sd zero,0(sp)
  sd t1,8(sp)
  RESTORE_CONTEXT

  jr ra
ENDPROC(do_scheduler)

ENTRY(ret_from_exception)
  /* TODO: */
  
  addi sp,sp,-24
  sd ra,0(sp)
  sd t0,8(sp)
  sd t1,16(sp)
  jal unlock_kernel
  ld ra,0(sp)
  ld t0,8(sp)
  ld t1,16(sp)
  addi sp,sp,24
  addi sp,sp,-16
  sd t1,8(sp)
  li t1,OFFSET_SIZE
  sd t1,0(sp)
  RESTORE_CONTEXT
  li t1,0xa22
  csrw sie,t1
  sret
ENDPROC(ret_from_exception)


